//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
#include "PDBToFASTA.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>

#include <fstream>
#include <iostream>

#include <zlib.h>
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
#define LENGTHLINE 5000
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
struct DataFasta {

  String* ide;
  String* sequence;
  String* file;

  DataFasta() {
    this->ide      = new String();
    this->sequence = new String();
    this->file     = new String();
  };

  ~DataFasta() {
    if (this->ide)      delete this->ide;
    if (this->sequence) delete this->sequence;
    if (this->file)     delete this->file;
  }

};

#define struct DataFasta DataFasta;
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
using namespace std;
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
namespace pdbFasta {

  bool AlphabeticSequences(TNode <DataFasta> *NodeT, void* ref) {
    DataFasta* seqA = NodeT->Info();
    DataFasta* seqB = ((TNode <DataFasta>*)ref)->Info();
    return (strcmp(seqA->sequence->Out(), seqB->sequence->Out()) > 0)?true:false;
  }

  bool EqualSequences(TNode <DataFasta> *NodeT, void* ref) {
    DataFasta* seqA   = NodeT->Info();
    DataFasta* seqB   = ((TNode <DataFasta>*)ref)->Info();
    bool       result = false;
    if (!strcmp(seqA->sequence->Out(), seqB->sequence->Out())) {
      String* temporal = seqB->ide->SubStr(1, (seqB->ide->Length()));
      seqA->ide->Add(",");
      seqA->ide->Add(temporal->Out());
      if (temporal) delete temporal;
      seqA->file->Add("\n");
      seqA->file->Add(seqB->file->Out());
      return true;
    }
    return false;
  }

};
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
PDBToFASTA::PDBToFASTA() {
  this->location            = new String();
  this->filenameOut         = new String();
  this->listInFilename      = new String();
  this->finalListInFilename = new String();
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
PDBToFASTA::PDBToFASTA(const PDBToFASTA& pdbToFASTA) {
  this->location            = new String(pdbToFASTA.location);
  this->filenameOut         = new String(pdbToFASTA.filenameOut);
  this->listInFilename      = new String(pdbToFASTA.listInFilename);
  this->finalListInFilename = new String(pdbToFASTA.finalListInFilename);
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
PDBToFASTA::PDBToFASTA(const PDBToFASTA* pdbToFASTA) {
  this->location            = new String(pdbToFASTA->location);
  this->filenameOut         = new String(pdbToFASTA->filenameOut);
  this->listInFilename      = new String(pdbToFASTA->listInFilename);
  this->finalListInFilename = new String(pdbToFASTA->finalListInFilename);
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
PDBToFASTA::~PDBToFASTA() {
  if (this->location)            delete this->location;
  if (this->filenameOut)         delete this->filenameOut;
  if (this->listInFilename)      delete this->listInFilename;
  if (this->finalListInFilename) delete this->finalListInFilename;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
void PDBToFASTA::Location(String* location) {
  if (location) {
    if (this->location) delete this->location;
    this->location = location;
  }
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
void PDBToFASTA::FilenameOut(String* filenameOut) {
  if (filenameOut) {
    if (this->filenameOut) delete this->filenameOut;
    this->filenameOut = filenameOut;
  }
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
void PDBToFASTA::ListInFilename(String* listInFilename) {
  if (listInFilename) {
    if (this->listInFilename) delete this->listInFilename;
    this->listInFilename = listInFilename;
  }
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
void PDBToFASTA::FinalListInFilename(String* finalListInFilename) {
  if (finalListInFilename) {
    if (this->finalListInFilename) delete this->finalListInFilename;
    this->finalListInFilename = finalListInFilename;
  }
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
String* PDBToFASTA::Location(void) {
  return this->location;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
String* PDBToFASTA::FilenameOut(void) {
  return this->filenameOut;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
String* PDBToFASTA::ListInFilename(void) {
  return this->listInFilename;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
String* PDBToFASTA::FinalListInFilename(void) {
  return this->finalListInFilename;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
bool PDBToFASTA::ConstructFileUsingList() {

  if (!this->location || !this->listInFilename || !this->filenameOut) return false;

  TListE <String> *listFiles = new TListE <String> ();
  TListE <String> *columns   = new TListE <String> ();
  
  char             buffer[LENGTHLINE + 1], buffer1[LENGTHLINE + 1];

  fstream          listFile;
  fstream          newListFile;
  fstream          outFile;
  
  char             completePath[strlen(this->location->Out()) + LENGTHLINE];

  listFile.open(this->listInFilename->Out(), fstream::in);
  while (!listFile.eof()) {
    listFile.getline(buffer, LENGTHLINE);
    String* nameFile = new String(buffer);
    nameFile->Trim();
    listFiles->Add(nameFile);
  }
  listFile.close();

  String* newFilename = new String(this->filenameOut->Out());
  newFilename->Add(".newList");

  outFile.open(this->filenameOut->Out(), fstream::out);
  newListFile.open(newFilename->Out(), fstream::out);

  if (newFilename) delete newFilename;

  cout << "Reading Files..." << endl;
  cout << "---------------------------------------------------------------------------------" << endl;

  listFiles->SetInitial();

  while (listFiles->GetNext()) {

    strcpy(completePath, this->location->Out());
    strcat(completePath, listFiles->GetCurrent()->Out());
    cout << "File: " << completePath << endl;

    gzFile pdbFile     = gzopen(completePath, "rb9");
    bool   condChain   = false;
    bool   skipChain   = false;
    int    posIniAA    = 4;

    String chainLetter, sequence, InitialPos;

    while (!gzeof(pdbFile)) {

      gzgets(pdbFile, buffer, LENGTHLINE);
      strcpy(buffer1, buffer);

      if (!strncmp(buffer, "SEQRES", strlen("SEQRES")) && !skipChain) {

        columns->Clear();
        this->SplitColumn(buffer1, columns, " ");

        if (!strcmp(columns->Get(1)->Out(), "1") && condChain) {
          skipChain = true;
          continue;
        }

        if (!strcmp(columns->Get(1)->Out(), "1") && (strlen(columns->Get(posIniAA)->Out()) == 3)) {
          condChain = true;
          chainLetter.In(columns->Get(2)->Out());
        }

        if (condChain) {
          for (int i = posIniAA; i < columns->Length(); i++)
            if (this->OneLetterAACode(columns->Get(i)) != ' ')
              sequence.Add(this->OneLetterAACode(columns->Get(i)));
        }

      }
      else if (condChain && !strncmp(buffer, "ATOM", strlen("ATOM"))) {
        String *line     = new String(buffer);
        String *temporal = line->SubStr(23, 26);
        InitialPos.In(temporal);
        if (temporal) delete temporal;
        if (line)     delete line;
        break;
      }

    }

    if (condChain && sequence.Length()) {

      String* pdbId = this->ExtractPDBCodeFromNamefile(listFiles->GetCurrent());

      outFile << ">" << pdbId->Out() << "_" << chainLetter.Out() << "_" << atoi(InitialPos.Out()) << endl;
      outFile << sequence.Out() << endl;

      cout << ">" << pdbId->Out() << "_" << chainLetter.Out() << "_" << atoi(InitialPos.Out()) << endl;
      cout << sequence.Out() << endl;

      newListFile << listFiles->GetCurrent()->Out() << endl;

      if (pdbId) delete pdbId;

    }

    gzclose(pdbFile);

  }

  outFile.close();
  newListFile.close();

  if (listFiles) delete listFiles;
  if (columns)   delete columns;

  return true;

}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
bool PDBToFASTA::ConstructFileWithoutDUsingList() {

  if (!this->location || !this->listInFilename || !this->filenameOut) return false;

  TListE <String>    *listFiles  = new TListE <String> ();
  TListE <String>    *columns    = new TListE <String> ();
  TListE <DataFasta> *listUnique = new TListE <DataFasta> ();

  char             buffer[LENGTHLINE + 1], buffer1[LENGTHLINE + 1];

  fstream          listFile;
  fstream          newListFile;
  fstream          outFile;

  char             completePath[strlen(this->location->Out()) + LENGTHLINE];

  listFile.open(this->listInFilename->Out(), fstream::in);
  while (!listFile.eof()) {
    listFile.getline(buffer, LENGTHLINE);
    String* nameFile = new String(buffer);
    nameFile->Trim();
    listFiles->Add(nameFile);
  }
  listFile.close();

  String* newFilename = new String(this->filenameOut->Out());
  newFilename->Add(".newList");

  outFile.open(this->filenameOut->Out(), fstream::out);
  newListFile.open(newFilename->Out(), fstream::out);

  if (newFilename) delete newFilename;

  cout << "Reading Files..." << endl;
  cout << "---------------------------------------------------------------------------------" << endl;

  listFiles->SetInitial();

  while (listFiles->GetNext()) {

    strcpy(completePath, this->location->Out());
    strcat(completePath, listFiles->GetCurrent()->Out());
    cout << "File: " << completePath << endl;

    gzFile pdbFile     = gzopen(completePath, "rb9");
    bool   condChain   = false;
    bool   skipChain   = false;
    int    posIniAA    = 4;

    String chainLetter, sequence, InitialPos;

    while (!gzeof(pdbFile)) {

      gzgets(pdbFile, buffer, LENGTHLINE);
      strcpy(buffer1, buffer);

      if (!strncmp(buffer, "SEQRES", strlen("SEQRES")) && !skipChain) {

        columns->Clear();
        this->SplitColumn(buffer1, columns, " ");

        if (!strcmp(columns->Get(1)->Out(), "1") && condChain) {
          skipChain = true;
          continue;
        }

        if (!strcmp(columns->Get(1)->Out(), "1") && (strlen(columns->Get(posIniAA)->Out()) == 3)) {
          condChain = true;
          chainLetter.In(columns->Get(2)->Out());
        }

        if (condChain) {

          bool condBreak = false;

          for (int i = posIniAA; i < columns->Length(); i++) {
            if (this->OneLetterAACode(columns->Get(i)) != ' ')
              sequence.Add(this->OneLetterAACode(columns->Get(i)));
            else if (this->OneLetterAACode(columns->Get(i)) == ' ' && columns->Get(i)->Length()) {
              condBreak = true;
              break;
            }
          }

          if (condBreak) {
            condChain = false;
            break;
          }
          
        }

      }
      else if (condChain && !strncmp(buffer, "ATOM", strlen("ATOM"))) {
        String *line     = new String(buffer);
        String *temporal = line->SubStr(22, 26);
        temporal->RTrim();
        InitialPos.In(temporal);
        if (temporal) delete temporal;
        if (line)     delete line;
        break;
      }

    }

    if (condChain && sequence.Length()) {

      String* pdbId = this->ExtractPDBCodeFromNamefile(listFiles->GetCurrent());

      DataFasta* data = new DataFasta();

      data->ide->Add(">");
      data->ide->Add(pdbId->Out());
      data->ide->Add("-");
      data->ide->Add(chainLetter.Out());
      data->ide->Add("_");
      data->ide->Add(InitialPos.Out());
      data->sequence->Add(sequence.Out());
      data->file->In(listFiles->GetCurrent()->Out());

      if (pdbId) delete pdbId;

      listUnique->Add(data);

    }

    gzclose(pdbFile);

  }

  cout << "Ordering the sequences...." << endl;
  listUnique->QuickSortList(pdbFasta::AlphabeticSequences, NULL);
  cout << "Reducing the list of equals sequences" << endl;
  listUnique->CompactListAfterOrdered(pdbFasta::EqualSequences, NULL);
  
  listUnique->SetInitial();

  while (listUnique->SetNext()) {

    outFile << listUnique->GetCurrent()->ide->Out() << endl;
    outFile << listUnique->GetCurrent()->sequence->Out() << endl;

    cout << listUnique->GetCurrent()->ide->Out() << endl;
    cout << listUnique->GetCurrent()->sequence->Out() << endl;
    
    newListFile << listUnique->GetCurrent()->file->Out() << endl;
    
  }

  outFile.close();
  newListFile.close();

  if (listFiles)  delete listFiles;
  if (columns)    delete columns;
  if (listUnique) delete listUnique;

  return true;

}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
bool PDBToFASTA::GetNameLigandPDBUsingList(String* locationDatabase) {

  if (!this->location || !this->listInFilename || !this->filenameOut) return false;

  TListE <String> *listFiles  = new TListE <String> ();
  char             buffer[LENGTHLINE + 1];

  fstream          listFile;
  fstream          outFile;
  gzFile           pdbFile;

  char             completePath[strlen(this->location->Out()) + LENGTHLINE];

  listFile.open(this->listInFilename->Out(), fstream::in);
  while (!listFile.eof()) {
    listFile.getline(buffer, LENGTHLINE);
    String* nameFile = new String(buffer);
    nameFile->RTrim();
    if (!nameFile->Length()) { 
      if (nameFile) delete nameFile;
      continue;
    }
    listFiles->Add(nameFile);
  }
  listFile.close();

  outFile.open(this->filenameOut->Out(), fstream::out);

  cout << "Reading Files..." << endl;
  cout << "---------------------------------------------------------------------------------" << endl;

  listFiles->SetInitial();

  while (listFiles->GetNext()) {

    sprintf(completePath, "%s%s\0", this->location->Out(), listFiles->GetCurrent()->Out());
    cout << "Ligand PDB File: " << completePath << endl;

    String* ligandID   = NULL;
    String* ligandName = NULL;
    String* temporal   = NULL;
    String* pdbId      = NULL;

    String  line;

    pdbFile = gzopen(completePath, "rb9");
    while (!gzeof(pdbFile)) {
      gzgets(pdbFile, buffer, LENGTHLINE);
      line.In(buffer);line.Trim();
      ligandID = line.SubStr(16, 19); ligandID->Trim();
      pdbId    = this->ExtractPDBCodeFromNamefile(listFiles->GetCurrent());
      break;
    }

    gzclose(pdbFile);

    strcpy(completePath, locationDatabase->Out());
    strcat(completePath, listFiles->GetCurrent()->Out());
    cout << "PDB File: " << completePath << endl;

    pdbFile = gzopen(completePath, "rb9");
    bool condHetATM = false;

    while (!gzeof(pdbFile)) {

      gzgets(pdbFile, buffer, LENGTHLINE);

      if      (strncmp(buffer, "HETNAM", strlen("HETNAM")) && !condHetATM) continue;
      else if (strncmp(buffer, "HETNAM", strlen("HETNAM")) &&  condHetATM) break;

      line.In(buffer);
      temporal = line.SubStr(11, 14); temporal->Trim();

      if (!strcmp(temporal->Out(), ligandID->Out())) {
        if (!condHetATM) {
          ligandName = line.SubStr(15, 70);
          ligandName->RTrim();
          condHetATM = true;
        }
        else {
          String* tmpLigandName = line.SubStr(15, 70);
          tmpLigandName->RTrim();
          ligandName->Add(tmpLigandName);
          if (tmpLigandName) delete tmpLigandName;
        }
      }

      if (temporal) delete temporal;

    }
    gzclose(pdbFile);

    outFile << pdbId->Out() << " " << ligandID->Out() << " " << ligandName->Out() << endl;
    cout << pdbId->Out() << " " << ligandID->Out() << " " << ligandName->Out() << endl;
    
    if (ligandID)   delete ligandID;
    if (ligandName) delete ligandName;
    if (pdbId)      delete pdbId;

  }
  
  outFile.close();

  if (listFiles)  delete listFiles;

  return true;

}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
bool PDBToFASTA::SplitColumn(char* buffer, TListE<String>* columns, char* characters) {

  if (!buffer || !columns || !characters) return false;

  char *part = strtok(buffer, characters);

  while (part) {
    String* cell = new String(part);
    cell->Trim();
    columns->Add(cell);
    part = strtok(NULL, characters);
  }

  return true;

}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
bool PDBToFASTA::SplitColumn(char* buffer, TListE<String>* columns, const char* characters) {

  if (!buffer || !columns || !characters) return false;

  char *part = strtok(buffer, characters);

  while (part) {
    String* cell = new String(part);
    cell->Trim();
    columns->Add(cell);
    part = strtok(NULL, characters);
  }

  return true;

}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
char PDBToFASTA::OneLetterAACode(String* code) {

  if      (!code)                return ' ';
  if      (code->Contain("ALA")) return 'A';
  else if (code->Contain("ARG")) return 'R';
  else if (code->Contain("ASN")) return 'N';
  else if (code->Contain("ASP")) return 'D';
  else if (code->Contain("ASX")) return 'B';
  else if (code->Contain("CYS")) return 'C';
  else if (code->Contain("GLU")) return 'E';
  else if (code->Contain("GLN")) return 'Q';
  else if (code->Contain("GLX")) return 'Z';
  else if (code->Contain("GLY")) return 'G';
  else if (code->Contain("HIS")) return 'H';
  else if (code->Contain("ILE")) return 'I';
  else if (code->Contain("LEU")) return 'L';
  else if (code->Contain("LYS")) return 'K';
  else if (code->Contain("MET")) return 'M';
  else if (code->Contain("PHE")) return 'F';
  else if (code->Contain("PRO")) return 'P';
  else if (code->Contain("SER")) return 'S';
  else if (code->Contain("THR")) return 'T';
  else if (code->Contain("TRP")) return 'W';
  else if (code->Contain("TYR")) return 'Y';
  else if (code->Contain("VAL")) return 'V';
  else                           return ' ';
  
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
String* PDBToFASTA::ExtractPDBCodeFromNamefile(String* nameFile) {
  if (!nameFile) return NULL;
  String* pdbId = new String();
  for (int i = 3; i < 7; i++)
    pdbId->Add(nameFile->Out(i));
  return pdbId;
}
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
//---------------------------------------------------------------------------------------------------------------------------------------------------------//
